{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import json\n",
    "import sys\n",
    "from torch import nn\n",
    "import torch\n",
    "from pytorch_transformers import *\n",
    "import importlib\n",
    "from tensorboardX import SummaryWriter\n",
    "import torch.nn.utils.rnn as rnn_utils\n",
    "import numpy as np\n",
    "import TNews_Loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 胶囊网络模型部分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "class sentence2capsules(nn.Module):\n",
    "    def __init__(self, in_dim_caps, out_num_caps, out_dim_caps, attn_dim=8, activation=None):\n",
    "        super(sentence2capsules, self).__init__()\n",
    "        self.in_dim_caps = in_dim_caps\n",
    "        self.out_num_caps = out_num_caps\n",
    "        self.out_dim_caps = out_dim_caps\n",
    "        self.activation = activation\n",
    "        self.attn_dim = attn_dim\n",
    "        self.q_mat = nn.Linear(in_dim_caps, out_num_caps*self.attn_dim)\n",
    "        self.k_mat = nn.Linear(in_dim_caps, out_num_caps*self.attn_dim)\n",
    "        self.v_mat = nn.Linear(in_dim_caps, out_num_caps*out_dim_caps)\n",
    "        \n",
    "    def forward(self, in_caps, sent_vec, attn_mask=None):\n",
    "        # in_caps = [batchsize, max_sent_len, in_dim_caps]\n",
    "        # sent_vec = [batchsize, in_dim_caps]\n",
    "        # attn_mask = [batchsize, max_sent_len]\n",
    "        batchsize, max_sent_len, _ = in_caps.shape\n",
    "        st_q = self.q_mat(st_vec).reshape(-1, 1, self.out_num_caps, self.attn_dim).transpose(1, 2)\n",
    "        caps_k = self.k_mat(in_caps).reshape(-1, max_sent_len, self.out_num_caps, self.attn_dim).transpose(1, 2)\n",
    "        scores = torch.matmul(st_q, caps_k.transpose(-1, -2))/np.sqrt(self.attn_dim) #[batchsize, out_num_caps, 1, max_sent_len]\n",
    "        attn = scores.transpose(0, 1) + attn_mask.unsqueeze(1) if attn_mask is not None else scores.transpose(0, 1)\n",
    "        attn_scores = torch.softmax(attn, dim=-1) #[out_num_caps, batchsize, 1, max_sent_len]\n",
    "        caps_v = self.v_mat(in_caps).reshape(-1, max_sent_len, self.out_num_caps, self.out_dim_caps).transpose(1, 2).transpose(0, 1)\n",
    "        out_caps = torch.matmul(attn_scores, caps_v).transpose(0, 1).squeeze(-2) # return [batchsize, out_num_caps, out_dim_caps]\n",
    "        if self.activation is not None:\n",
    "            out_caps = self.activation(out_caps)\n",
    "        return out_caps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "st_vec = torch.randn(5, 10)\n",
    "wd_vecs = torch.randn(5, 12, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "s2caps = sentence2capsules(10, 6, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class caps2class(nn.Module):\n",
    "    def __init__(self, in_num_caps, in_dim_caps, out_num_caps, out_dim_caps, routings):\n",
    "        super(caps2class, self).__init__()\n",
    "        self.in_num_caps = in_num_caps\n",
    "        self.in_dim_caps = in_dim_caps\n",
    "        self.out_num_caps = out_num_caps\n",
    "        self.out_dim_caps = out_dim_caps\n",
    "        self.routings = routings\n",
    "        \n",
    "    def forward(self, in_caps, st_vec):\n",
    "\n",
    "        \n",
    "        \n",
    "    def dynamic_routing(self, x_hat):\n",
    "        b = Variable(torch.zeros(x_hat.size(0), self.out_num_caps, self.in_num_caps)).cuda()\n",
    "        assert self.routings > 0, 'The \\'routings\\' should be > 0.'\n",
    "        for i in range(self.routings):\n",
    "            c = F.softmax(b, dim=1)\n",
    "            if i == self.routings - 1:\n",
    "                outputs = squash(torch.sum(c[:, :, :, None] * x_hat, dim=-2, keepdim=True))\n",
    "            else:\n",
    "                outputs = squash(torch.sum(c[:, :, :, None] * x_hat_detached, dim=-2, keepdim=True))\n",
    "                b = b + torch.sum(outputs * x_hat_detached, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def squash(inputs, axis=-1):\n",
    "    norm = torch.norm(inputs, p=2, dim=axis, keepdim=True)\n",
    "    scale = norm**2 / (1 + norm**2) / (norm + 1e-8)\n",
    "    return scale * inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 主函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "xt = XLNetTokenizer.from_pretrained(\"./XLBaseModel/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "xm = XLNetModel.from_pretrained(\"./XLBaseModel/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_file = './THUCnews/cnews.test.txt'\n",
    "val_file = './THUCnews/cnews.val.txt'\n",
    "train_file = './THUCnews/cnews.train.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_data = TNews_Loader.load_data(\"./THUCnews/cnews.val.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_reader = TNews_Loader.THUReader(val_data, 10, xt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "for x, y, l in val_reader.iter():\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "rst = xm(torch.tensor([x[9]]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
